package com.example.soloadsso;

import jcifs.Config;
import jcifs.UniAddress;
import jcifs.ntlmssp.Type1Message;
import jcifs.ntlmssp.Type2Message;
import jcifs.ntlmssp.Type3Message;
import jcifs.smb.NtlmPasswordAuthentication;
import jcifs.smb.SmbAuthException;
import jcifs.smb.SmbSession;
import lombok.extern.slf4j.Slf4j;
import org.noear.solon.annotation.Component;
import org.noear.solon.annotation.Init;
import org.noear.solon.annotation.Inject;
import org.noear.solon.core.handle.Context;
import org.noear.solon.core.handle.Filter;
import org.noear.solon.core.handle.FilterChain;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

@Component
@Slf4j
public class AdSsoFilter implements Filter {
    @Inject("${adsso.domainServer}")
    private String domainServer;
    @Inject("${adsso.username}")
    private String username;
    @Inject("${adsso.password}")
    private String password;

    @Init
    public void init() {
        Config.setProperty("jcifs.http.domainController", domainServer);
        Config.setProperty("jcifs.smb.client.username", username);
        Config.setProperty("jcifs.smb.client.password", password);
        Config.setProperty("jcifs.smb.client.useExtendedSecurity", "false");
        Config.setProperty("jcifs.smb.lmCompatibility", "0");
    }

    public String getUserName() {
        try {
            NtlmPasswordAuthentication ntlm = negotiate();
            if (ntlm == null) {
                return null;
            }
            return ntlm.getUsername();
        } catch (Throwable e) {
            log.error("域环境自动登录失败", e);
        }
        return null;
    }

    private NtlmPasswordAuthentication negotiate() throws IOException {
        Context ctx = Context.current();
        UniAddress dc;
        String msg = ctx.header("Authorization");
        NtlmPasswordAuthentication ntlm;
        if (msg != null && (msg.startsWith("NTLM "))) {
            byte[] challenge;
            dc = UniAddress.getByName(domainServer, true);
            challenge = SmbSession.getChallenge(dc);
            if ((ntlm = authenticate(challenge)) == null) {
                return null;
            }
            ctx.sessionRemove("NtlmHttpChal");
            try {
                SmbSession.logon(dc, ntlm);
            } catch (SmbAuthException sae) {
                if (sae.getNtStatus() == sae.NT_STATUS_ACCESS_VIOLATION) {
                    ctx.sessionRemove("NtlmHttpAuth");
                }
                ctx.headerSet("WWW-Authenticate", "NTLM");
                ctx.status(HttpServletResponse.SC_UNAUTHORIZED);
                ctx.contentLength(0);
                ctx.flush();
                return null;
            }
            ctx.sessionSet("NtlmHttpAuth", ntlm);
        } else {
            if ((ntlm = (NtlmPasswordAuthentication) ctx.session("NtlmHttpAuth")) == null) {
                ctx.headerSet("WWW-Authenticate", "NTLM");
                ctx.status(HttpServletResponse.SC_UNAUTHORIZED);
                ctx.contentLength(0);
                ctx.flush();
                return null;
            }
        }
        return ntlm;
    }

    private NtlmPasswordAuthentication authenticate(byte[] challenge) throws IOException {
        Context ctx = getContext();
        String msg = ctx.header("Authorization");
        if (msg != null && msg.startsWith("NTLM ")) {
            byte[] src = jcifs.util.Base64.decode(msg.substring(5));
            if (src[8] == 1) {
                Type1Message type1 = new Type1Message(src);
                Type2Message type2 = new Type2Message(type1, challenge, null);
                msg = jcifs.util.Base64.encode(type2.toByteArray());
                ctx.headerSet("WWW-Authenticate", "NTLM " + msg);
            } else if (src[8] == 3) {
                Type3Message type3 = new Type3Message(src);
                byte[] lmResponse = type3.getLMResponse();
                if (lmResponse == null) lmResponse = new byte[0];
                byte[] ntResponse = type3.getNTResponse();
                if (ntResponse == null) ntResponse = new byte[0];
                return new NtlmPasswordAuthentication(type3.getDomain(),
                        type3.getUser(), challenge, lmResponse, ntResponse);
            }
        } else {
            ctx.headerSet("WWW-Authenticate", "NTLM");
        }
        ctx.status(HttpServletResponse.SC_UNAUTHORIZED);
        ctx.contentLength(0);
        ctx.flush();
        return null;
    }

    private Context getContext() {
        Context ctx = Context.current();
        if (ctx == null) {
            throw new RuntimeException("非web环境下无法获取Context");
        }
        return ctx;
    }

    @Override
    public void doFilter(Context ctx, FilterChain chain) {
        try {
            if (ctx.path().equals("/adsso")) {
                String username = getUserName();
                if (username == null || username.isEmpty()) {
                    chain.doFilter(ctx);
                    return;
                }
                //后续的用户登录操作，将token写入cookies并跳转到首页或者其他的登录方式
                ctx.render("域用户自动登录成功：" + username);
            } else {
                chain.doFilter(ctx);
            }
        } catch (Throwable e) {
            log.error("AdSsoFilterError", e);
        }
    }
}
